{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}

-- | Module: Data.Vector.Sized.Sort
-- Copyright: (c) 2018 Koz Ross
-- Maintainer: Koz Ross <koz.ross@retro-freedom.nz>
-- Stability: Experimental
-- Portability: Non-portable ('ScopedTypeVariables', 'TypeFamilies')
--
-- A range of implementations of sorting algorithms, based on the
-- vector-algorithms package. These return sorted copies; see
-- "Data.Vector.Sized.Mutable.Sort" for mutable versions of each of these.
module Data.Vector.Sized.Sort
  ( Lexicographic
  , Comparison
  , americanFlagSort
  , americanFlagSortBy
  , heapSort
  , heapSortBy
  , heapSortOn
  , insertionSort
  , insertionSortBy
  , insertionSortOn
  , introSort
  , introSortBy
  , introSortOn
  , mergeSort
  , mergeSortBy
  , mergeSortOn
  , timSort
  , timSortBy
  , timSortOn
  , discriminantSort
  , discriminantSortOn
  ) where

import Control.Monad.Primitive
import Control.Monad.ST
import Data.Coerce
import Data.Discrimination
import Data.Function
import Data.Functor.Contravariant
import Data.Vector.Algorithms.AmericanFlag (Lexicographic)
import GHC.TypeNats

import qualified Data.Vector.Algorithms.AmericanFlag as AF
import qualified Data.Vector.Algorithms.Heap as H
import qualified Data.Vector.Algorithms.Insertion as I
import qualified Data.Vector.Algorithms.Intro as IS
import qualified Data.Vector.Algorithms.Merge as M
import qualified Data.Vector.Algorithms.Tim as T
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable.Sized as VGMS
import qualified Data.Vector.Generic.Sized as VGS
import qualified Data.Vector.Mutable.Sized.Sort as VMSS

-- | Returns a sorted copy using American flag sort. For details, see
-- Data.Vector.Algorithms.AmericanFlag.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE americanFlagSort #-}
americanFlagSort ::
     (KnownNat n, Lexicographic a, Ord a, VG.Vector v a)
  => VGS.Vector v n a
  -> VGS.Vector v n a
americanFlagSort vec = runST (unpackMutable AF.sort vec)

-- | Fine-grained version of 'americanFlagSort'.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE americanFlagSortBy #-}
americanFlagSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a -- ^ Comparison for insertion sort fallback
  -> (a -> Int -> Bool) -- ^ Is a stripe complete?
  -> Int -- ^ Number of buckets necessary.
  -> (Int -> a -> Int) -- ^ Big-endian radix function.
  -> VGS.Vector v n a
  -> VGS.Vector v n a
americanFlagSortBy comp strip buckets radix vec =
  runST (unpackMutable (AF.sortBy (coerce comp) strip buckets radix) vec)

-- | Returns a sorted copy using heapsort with a quarternary heap. For details,
-- see Data.Vector.Algorithms.Heap.
-- 
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE heapSort #-}
heapSort ::
     (KnownNat n, Ord a, VG.Vector v a) => VGS.Vector v n a -> VGS.Vector v n a
heapSort vec = runST (unpackMutable H.sort vec)

-- | As 'heapSort', but with a custom 'Comparison'.
{-# INLINE heapSortBy #-}
heapSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> VGS.Vector v n a
heapSortBy comp vec = runST (unpackMutable (H.sortBy (coerce comp)) vec)

-- | As 'heapSort', using a \'projection\' to determine ordering.
{-# INLINE heapSortOn #-}
heapSortOn ::
     (KnownNat n, Ord b, VG.Vector v a)
  => (a -> b)
  -> VGS.Vector v n a
  -> VGS.Vector v n a
heapSortOn f vec = runST (unpackMutable (H.sortBy go) vec)
  where
    go = coerce (compare `on` f)

-- | Returns a sorted copy using insertion sort. For details, see
-- Data.Vector.Algorithms.Insertion.
--
-- __Complexity:__ \(O(n^{2})\)
{-# INLINE insertionSort #-}
insertionSort ::
     (KnownNat n, Ord a, VG.Vector v a) => VGS.Vector v n a -> VGS.Vector v n a
insertionSort vec = runST (unpackMutable I.sort vec)

-- | As 'insertionSort', but with a custom Comparison.
{-# INLINE insertionSortBy #-}
insertionSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> VGS.Vector v n a
insertionSortBy comp vec = runST (unpackMutable (I.sortBy (coerce comp)) vec)

-- | As 'insertionSort', using a \'projection\' to determine ordering.
{-# INLINE insertionSortOn #-}
insertionSortOn ::
     (KnownNat n, Ord b, VG.Vector v a)
  => (a -> b)
  -> VGS.Vector v n a
  -> VGS.Vector v n a
insertionSortOn f vec = runST (unpackMutable (I.sortBy go) vec)
  where
    go = coerce (compare `on` f)

-- | Returns a sorted copy using introsort. For details, see
-- Data.Vector.Algorithms.Intro.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE introSort #-}
introSort ::
     (KnownNat n, Ord a, VG.Vector v a) => VGS.Vector v n a -> VGS.Vector v n a
introSort vec = runST (unpackMutable IS.sort vec)

-- | As 'introSort', but with a custom 'Comparison'.
{-# INLINE introSortBy #-}
introSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> VGS.Vector v n a
introSortBy comp vec = runST (unpackMutable (IS.sortBy (coerce comp)) vec)

-- | As 'introSort', with a \'projection\' to determine ordering.
{-# INLINE introSortOn #-}
introSortOn ::
     (KnownNat n, Ord b, VG.Vector v a)
  => (a -> b)
  -> VGS.Vector v n a
  -> VGS.Vector v n a
introSortOn f vec = runST (unpackMutable (IS.sortBy go) vec)
  where
    go = coerce (compare `on` f)

-- | Returns a sorted copy using top-down mergesort. For details, see
-- Data.Vector.Algorithms.Merge.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE mergeSort #-}
mergeSort ::
     (KnownNat n, Ord a, VG.Vector v a) => VGS.Vector v n a -> VGS.Vector v n a
mergeSort vec = runST (unpackMutable M.sort vec)

-- | As 'mergeSort', but with a custom Comparison.
{-# INLINE mergeSortBy #-}
mergeSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> VGS.Vector v n a
mergeSortBy comp vec = runST (unpackMutable (M.sortBy (coerce comp)) vec)

-- | As 'mergeSort', using a \'projection\' to determine ordering.
{-# INLINE mergeSortOn #-}
mergeSortOn ::
     (KnownNat n, Ord b, VG.Vector v a)
  => (a -> b)
  -> VGS.Vector v n a
  -> VGS.Vector v n a
mergeSortOn f vec = runST (unpackMutable (M.sortBy go) vec)
  where
    go = coerce (compare `on` f)

-- | Returns a sorted copy using Timsort. For details, see
-- Data.Vector.Algorithms.Tim.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE timSort #-}
timSort ::
     (KnownNat n, Ord a, VG.Vector v a) => VGS.Vector v n a -> VGS.Vector v n a
timSort vec = runST (unpackMutable T.sort vec)

-- | As 'timSort', but with a custom 'Comparison'.
{-# INLINE timSortBy #-}
timSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> VGS.Vector v n a
timSortBy comp vec = runST (unpackMutable (T.sortBy (coerce comp)) vec)

-- | As 'timSort', using a \'projection\' to determine ordering.
{-# INLINE timSortOn #-}
timSortOn ::
     (KnownNat n, Ord b, VG.Vector v a)
  => (a -> b)
  -> VGS.Vector v n a
  -> VGS.Vector v n a
timSortOn f vec = runST (unpackMutable (T.sortBy go) vec)
  where
    go = coerce (compare `on` f)

-- | Returns a sorted copy using a stable, ordered discriminator. For details,
-- see 'Data.Discrimination'.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE discriminantSort #-}
discriminantSort ::
     (KnownNat n, Sorting a, VG.Vector v a)
  => VGS.Vector v n a
  -> VGS.Vector v n a
discriminantSort vec =
  runST
    (do v <- VGMS.unsafeNew
        VGS.copy v vec
        VMSS.discriminantSort v
        VGS.unsafeFreeze v)

-- | As 'discriminantSort', using a \'projection\' to determine the
-- discriminator.
{-# INLINE discriminantSortOn #-}
discriminantSortOn ::
     (KnownNat n, Sorting b, VG.Vector v a)
  => (a -> b)
  -> VGS.Vector v n a
  -> VGS.Vector v n a
discriminantSortOn f vec =
  runST
    (do v <- VGMS.unsafeNew
        VGS.copy v vec
        VMSS.discriminantSortOn f v
        VGS.unsafeFreeze v)

-- Helper to wrap 'unpacking' into mutable arrays
{-# INLINE unpackMutable #-}
unpackMutable ::
     (KnownNat n, PrimMonad m, VG.Vector v a)
  => (VG.Mutable v (PrimState m) a -> m ())
  -> VGS.Vector v n a
  -> m (VGS.Vector v n a)
unpackMutable f vec = do
  v <- VGMS.unsafeNew
  VGS.copy v vec
  let w = VGMS.fromSized v
  f w
  VGS.unsafeFreeze v
